#pragma once
#include <eigen3/Eigen/Dense>
#include <eigen3/Eigen/Core>
#include <fstream>
#include <random>
#include <chrono>
#include <cmath>
#include "base/defines.h"

class DbgFrame;

namespace Filtering
{
    class State
    {
    public:
        State() : x_(0), y_(0), theta_(0)
        {
        }

        double x() const { return x_; }
        double y() const { return y_; }
        double theta() const { return theta_; }

        double &x() { return x_; }
        double &y() { return y_; }
        double &theta() { return theta_; }

    private:
        double x_;
        double y_;
        double theta_;
    };

    class Control
    {
    public:
        Control() : v_(0), dt_(0)
        {
        }

        double v() const { return v_; }
        double dt() const { return dt_; }

        double &v() { return v_; }
        double &dt() { return dt_; }

    private:
        double v_;
        double dt_;
    };

    class Measurement
    {
    public:
        Measurement() : x_(0), y_(0), theta_(0)
        {
        }

        Measurement(double x, double y, double theta) : x_(x), y_(y), theta_(theta)
        {
        }

        double x() const { return x_; }
        double y() const { return y_; }
        double theta() const { return theta_; }

        double &x() { return x_; }
        double &y() { return y_; }
        double& theta() { return theta_; }
    private:
        double x_;
        double y_;
        double theta_;
    };

    class ActionModel
    {
    public:
        ActionModel()
        {
            cov_ << 0.5, 0.0, 0.0,
                0.0, 0.5, 0.0,
                0.0, 0.0, 0.1;
        }

        State transition(const State &x, const Control &u)
        {
            State x_t;
            x_t.x() = x.x() + std::cos(x_t.theta()) * u.v() * u.dt();
            x_t.y() = x.y() + std::sin(x_t.theta()) * u.v() * u.dt();
            x_t.theta() = x.theta();
            return x_t;
        }

        void updateJacobians(const State &x, const Control &u)
        {
            F_.setZero();
            F_(0, 0) = 1.0;
            F_(0, 1) = 0.0;
            F_(0, 2) = -std::sin(x.theta()) * u.v() * u.dt();

            F_(1, 0) = 0.0;
            F_(1, 1) = 1.0;
            F_(1, 2) = std::cos(x.theta()) * u.v() * u.dt();

            F_(2, 0) = 0.0;
            F_(2, 1) = 0.0;
            F_(2, 2) = 1.0;
        }

        const Eigen::Matrix3d &F() const
        {
            return F_;
        }

        const Eigen::Matrix3d &getCovariance() const
        {
            return cov_;
        }

    private:
        Eigen::Matrix3d F_;
        Eigen::Matrix3d cov_;
    };

    class MeasureModel
    {
    public:
        MeasureModel()
        {
            cov_ << 5.0, 0.0, 0.0,
                0.0, 5.0, 0.0,
                0.0, 0.0, 0.1;
        }

        Measurement measure(const State &x) const
        {
            return Measurement(x.x(), x.y(), x.theta());
        }

        void updateJacobians(const State &x)
        {
            H_.setZero();
            H_(0, 0) = 1.0;
            H_(0, 1) = 0.0;
            H_(0, 2) = 0.0;

            H_(1, 0) = 0.0;
            H_(1, 1) = 1.0;
            H_(1, 2) = 0.0;

            H_(2, 0) = 0.0;
            H_(2, 1) = 0.0;
            H_(2, 2) = 1.0;
        }

        const Eigen::Matrix3d &H() const
        {
            return H_;
        }

        const Eigen::Matrix3d &getCovariance() const
        {
            return cov_;
        }

    private:
        Eigen::Matrix3d H_;
        Eigen::Matrix3d cov_;
    };

    class ExtendedKalmanFilter
    {
    public:

        void Init(GeoPoint x0, Eigen::Matrix3d p0)
        {
            x_.x() = x0.x;
            x_.y() = x0.y;
            x_.theta() = x0.bearing;
            P_ = p0;
            last_p = x0;
        }

        GeoPoint Predict(const GeoPoint& gps_point);

        bool Update(const GeoPoint& gps_point);

        void Print(DbgFrame* dbg_frame);

    private:
        const State &predict(ActionModel &s, const Control &u)
        {
            s.updateJacobians(x_, u);
            x_ = s.transition(x_, u);
            P_ = (s.F() * P_ * s.F().transpose()) + s.getCovariance();
            return x_;
        }

        const State &update(MeasureModel &m, const Measurement &z)
        {
            m.updateJacobians(x_);
            Eigen::Matrix3d S = (m.H() * P_ * m.H().transpose()) + m.getCovariance();
            Eigen::Matrix3d K = P_ * m.H().transpose() * S.inverse();

            Eigen::Vector3d delta;
            delta.x() = z.x() - m.measure(x_).x();
            delta.y() = z.y() - m.measure(x_).y();
            delta.z() = z.theta() - m.measure(x_).theta();
            Eigen::Vector3d rise = K * delta;

            x_.x() = x_.x() + rise.x();
            x_.y() = x_.y() + rise.y();
            x_.theta() = x_.theta() + rise.z();

            P_ -= K * m.H() * P_;
            return x_;
        }

    private:
        GeoPoint last_p;
        ActionModel sys_;
        MeasureModel pm_;
        Eigen::Matrix3d P_;
        State x_;
        std::vector<GeoPoint> predict_points;
    };
} // namespace EKF